package logs

import (
	cmap "github.com/orcaman/concurrent-map/v2"
	log "github.com/sirupsen/logrus"
	"sync"
	"sync/atomic"
	"time"
)

const (
	g_max_query_log = 100 * 10
	g_max_all_log   = 100 * 10
)

var (
	g_qps_task       *QPSTask
	g_http_dns_task  *HttpDNSTask
	g_query_log_task *QueryLogTask

	g_package_count_task *PackageCount
	g_all_log_task       *AllLog

	g_task_interval = time.Second * 30

	g_alllog_running atomic.Bool
)

type StatisticInterface interface {
	AddData(data interface{}) error
	SendData() error
}

type HostCount struct {
	Host      string    `json:"host"`
	Total     uint64    `json:"total"`
	Published time.Time `json:"timestamp"`
}
type PackageCount struct {
	index string
	HostCount
}
type AllLogData struct {
	Host   string `json:"host"`
	Remote string `json:"query_machine"`

	Domain    string    `json:"domain"`
	QueryType string    `json:"query_type"`
	Data      string    `json:"data"`
	Published time.Time `json:"timestamp"`
}
type AllLog struct {
	index string
	data  []interface{}

	rw_mutex *sync.RWMutex
	log_data atomic.Value
}
type QPSTask struct {
	index string
	data  cmap.ConcurrentMap[string, *QPSData]
}
type HttpDNSTask struct {
	index string
	data  cmap.ConcurrentMap[string, *HttpDNSData]
}
type QueryLogTask struct {
	index    string
	data     []interface{}
	rw_mutex *sync.RWMutex
	log_data atomic.Value
}

func init() {
	g_qps_task = NewQPSTask("root-qps")
	//g_http_dns_task = NewHttpDNSTask("http-dns")
	g_query_log_task = NewQueryLogTask("root-query-log")

	g_package_count_task = NewPackageCount("root-package-count")
	g_all_log_task = NewAllLog("root-log")

	g_alllog_running.Store(false)
}

func NewQPSTask(index string) *QPSTask {
	return &QPSTask{
		index: index,
		data:  cmap.New[*QPSData](),
	}
}

func NewHttpDNSTask(index string) *HttpDNSTask {
	return &HttpDNSTask{
		index: index,
		data:  cmap.New[*HttpDNSData](),
	}
}

func NewQueryLogTask(index string) *QueryLogTask {
	return &QueryLogTask{
		index:    index,
		rw_mutex: &sync.RWMutex{},
		data:     make([]interface{}, 0),
	}
}

func NewPackageCount(index string) *PackageCount {
	return &PackageCount{
		index: index,
		HostCount: HostCount{
			Host:  g_local_addr,
			Total: 0,
		},
	}
}

func NewAllLog(index string) *AllLog {
	return &AllLog{
		index:    index,
		rw_mutex: &sync.RWMutex{},
		data:     make([]interface{}, 0),
	}
}

func AddQpsData(domain string, is_udp bool) error {
	return g_qps_task.AddData(domain, is_udp)
}

func (q *QPSTask) AddData(domain string, is_udp bool) error {
	if ok := q.data.Has(domain); !ok {
		q.data.Set(domain, &QPSData{
			Host:      g_local_addr,
			Domain:    domain,
			Total:     0,
			Tcp:       0,
			Udp:       0,
			Published: time.Now().Round(time.Second),
		})
	}

	data, _ := q.data.Get(domain)
	if is_udp {
		data.Udp += 1
	} else {
		data.Tcp += 1
	}
	data.Total += 1

	return nil
}

func (q *QPSTask) SendData() error {
	if q.data.Count() == 0 {
		return nil
	}

	var data = make([]interface{}, 0)
	for _, v := range q.data.Items() {
		data = append(data, v)
	}

	//log.Errorf("send qps data to es, count: %d, value: %v", len(data), data)
	//init q.data
	q.data.Clear()

	//send data
	return SendDataToEs(q.index, data)
}

func AddHttpDnsData(domain, query_type string, is_aes, is_des, is_https bool) error {
	return g_http_dns_task.AddData(domain, query_type, is_aes, is_des, is_https)
}

func (h *HttpDNSTask) AddData(domain, query_type string, is_aes, is_des, is_https bool) error {
	if ok := h.data.Has(domain); !ok {
		h.data.Set(domain, &HttpDNSData{
			Host:      g_local_addr,
			Domain:    domain,
			QueryType: query_type,
			Published: time.Now().Round(time.Second),
		})
	}

	data, _ := h.data.Get(domain)
	if is_aes {
		data.Aes += 1
	}
	if is_des {
		data.Des += 1
	}
	if is_https {
		data.Https += 1
	}
	data.Total += 1

	return nil
}

func (h *HttpDNSTask) SendData() error {
	if h.data.Count() == 0 {
		return nil
	}

	var data = make([]interface{}, 0)
	for _, v := range h.data.Items() {
		data = append(data, v)
	}

	//init q.data
	h.data.Clear()

	//send data
	return SendDataToEs(h.index, data)
}

func AddQueryLogData(data QueryLog) {
	data.Host = g_local_addr
	g_query_log_task.AddData(&data)
}

func (q *QueryLogTask) AddData(data interface{}) {
	q.rw_mutex.Lock()
	defer q.rw_mutex.Unlock()

	if len(q.data) >= g_max_query_log {
		return
	}
	q.data = append(q.data, data)
	return
}

func (q *QueryLogTask) SendData() error {
	q.rw_mutex.Lock()
	defer q.rw_mutex.Unlock()

	if len(q.data) == 0 {
		return nil
	}

	q.log_data.Store(q.data)
	q.data = make([]interface{}, 0)

	return SendDataToEs(q.index, q.log_data.Load().([]interface{}))
}

func AddPackageCount() error {
	return g_package_count_task.AddData()
}

func (p *PackageCount) AddData() error {
	atomic.AddUint64(&p.Total, 1)
	return nil
}

func (p *PackageCount) SendData() error {
	total_count := atomic.LoadUint64(&p.Total)
	//if total_count == 0 {
	//	return nil
	//}
	count := HostCount{
		Host:      g_local_addr,
		Total:     total_count,
		Published: time.Now().Round(time.Second),
	}
	atomic.StoreUint64(&p.Total, 0)

	var data = make([]interface{}, 0)
	data = append(data, count)

	return SendDataToEs(p.index, data)
}

func AddAllLogData(data, remote, domain, query_type string) {
	if len(g_all_log_task.data) >= g_max_all_log {
		return
	}

	log_data := &AllLogData{
		Host:      g_local_addr,
		Remote:    remote,
		Data:      data,
		Domain:    domain,
		QueryType: query_type,
		Published: time.Now().Round(time.Second),
	}

	g_all_log_task.AddData(log_data)
}

func (a *AllLog) AddData(data *AllLogData) {
	a.rw_mutex.Lock()
	defer a.rw_mutex.Unlock()

	a.data = append(a.data, data)
	//a.log_data.Store(a.data)
	//log.Errorf("add all log data: %v", a.log_data.Load())
}

func (a *AllLog) SendData() error {
	a.rw_mutex.Lock()
	defer a.rw_mutex.Unlock()

	a.log_data.Store(a.data)

	a.data = nil

	return SendDataToEs(a.index, a.log_data.Load().([]interface{}))
}

func Run(ip string, es_addr, es_pwd, http_ca string) {
	//init
	var err error
	g_es_client.Es, err = g_es_client.initEsClient(es_addr, es_pwd, http_ca)
	if err != nil {
		log.Errorf("init es client error: %v", err)
		return
	} else {
		log.Tracef("init es client success")
	}

	//init test data
	//InitData()

	g_local_addr = ip
	for true {
		func() {
			defer func() {
				if err := recover(); err != nil {
					log.Errorf("statistic run error: %v", err)
				}
			}()

			wait := sync.WaitGroup{}
			wait.Add(4)

			go func() {
				err := g_qps_task.SendData()
				if err != nil {
					log.Errorf("statistic qps error: %v", err)
				}
				wait.Done()
			}()

			//go func() {
			//	err := g_http_dns_task.SendData()
			//	if err != nil {
			//		log.Errorf("statistic http_dns error: %v", err)
			//	}
			//	wait.Done()
			//}()

			go func() {
				err := g_query_log_task.SendData()
				if err != nil {
					log.Errorf("statistic query_log error: %v", err)
				}
				wait.Done()
			}()

			go func() {
				err := g_package_count_task.SendData()
				if err != nil {
					log.Errorf("statistic package_count error: %v", err)
				}
				wait.Done()
			}()

			go func() {
				err := g_all_log_task.SendData()
				if err != nil {
					log.Errorf("statistic all_log error: %v", err)
				}
				wait.Done()
			}()

			wait.Wait()
		}()

		time.Sleep(g_task_interval)
	}
}

// set task running status
func SetQueryLogStatus(status bool) {
	g_alllog_running.Store(status)
}

// get task running status
func GetQueryLogStatus() bool {
	return g_alllog_running.Load()
}
